import streamlit as st
import pandas as pd
from io import BytesIO
import sys
from contextlib import redirect_stdout, redirect_stderr
import io
import traceback
from code_interpreter import CodeInterpreter
import plotly.express as px
import plotly.graph_objects as go
import numpy as np
import ast

# 初始化代码解释器
if 'interpreter' not in st.session_state:
    st.session_state.interpreter = CodeInterpreter()

# 设置页面标题
st.title("Excel 处理工具 & 代码解释器")

# 上传文件
uploaded_file = st.file_uploader("上传 Excel 文件", type=["xlsx", "xls"])

# 全局变量用于存储DataFrame
if 'df' not in st.session_state:
    st.session_state.df = None

if uploaded_file is not None:
    # 读取 Excel 文件
    st.session_state.df = pd.read_excel(uploaded_file)

    # 显示原始数据
    st.subheader("原始数据")
    st.write(st.session_state.df)

    # 显示可用的列名
    st.sidebar.subheader("可用的数据列:")
    st.sidebar.write(", ".join(st.session_state.df.columns))

    # 代码编辑器
    st.subheader("Python 代码编辑器")
    default_code = """# 在这里编写处理Excel的Python代码
# 使用变量 df 来引用上传的Excel数据
# 支持的模块: pandas(pd), numpy(np), math, random, json, re, time, copy

# 图表生成示例 (请确保引号完整配对):
'''
# 柱状图示例
fig = px.bar(df, 
    x='姓名',        # x轴列名
    y='分数1',       # y轴列名
    title='成绩柱状图'  # 图表标题
)

# 折线图示例
fig = px.line(df,
    x='日期',        # x轴列名
    y='销售额',      # y轴列名
    title='销售趋势'  # 图表标题
)

# 散点图示例
fig = px.scatter(df,
    x='身高',        # x轴列名
    y='体重',        # y轴列名
    title='身高体重分布'  # 图表标题
)

# 饼图示例
fig = px.pie(df,
    values='数量',    # 数值列名
    names='类别',     # 分类列名
    title='类别分布'   # 图表标题
)
'''
"""
    code = st.text_area("编辑代码", value=default_code, height=300, key="code_editor")

    # 创建一个包含运行按钮和生成图表按钮的列布局
    col1, col2, col3 = st.columns([4, 1, 1])
    with col2:
        run_button = st.button("运行代码", type="primary")
    with col3:
        chart_button = st.button("生成图表", type="secondary")

    if run_button and code:
        st.subheader("执行结果")

        # 使用代码解释器执行代码
        result = st.session_state.interpreter.execute_code(code, st.session_state.df)

        if result['success']:
            # 更新DataFrame
            st.session_state.df = result['df']

            # 显示代码输出
            if result['output']:
                st.text("代码输出:")
                st.code(result['output'])

            # 显示处理后的数据
            st.subheader("处理后的数据")
            st.write(st.session_state.df)

            # 提供下载处理后的文件选项
            output_excel = BytesIO()
            with pd.ExcelWriter(output_excel, engine='openpyxl') as writer:
                st.session_state.df.to_excel(writer, index=False)
            output_excel.seek(0)

            st.download_button(
                label="下载处理后的 Excel 文件",
                data=output_excel,
                file_name="processed_excel.xlsx",
                mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
            )
        else:
            st.error("执行代码时出错:")
            st.code(result['error'])

    elif chart_button and code:
        st.subheader("图表生成")

        try:
            # 检查代码语法
            ast.parse(code)

            # 准备图表生成环境
            chart_env = {
                'pd': pd,
                'df': st.session_state.df,
                'px': px,
                'go': go,
                'np': np,
                'fig': None
            }

            # 执行代码
            exec(code, chart_env)

            # 检查是否生成了图表
            if 'fig' in chart_env and chart_env['fig'] is not None:
                st.plotly_chart(chart_env['fig'], use_container_width=True)
            else:
                st.warning("""未检测到图表对象。请确保代码中创建了 'fig' 变量。

示例代码：
```python
fig = px.bar(df, x='列名1', y='列名2', title='图表标题')
```

可用的列名：
""" + ", ".join(st.session_state.df.columns))

                # 显示简单的数据统计
                st.subheader("数据统计:")
                st.write(st.session_state.df.describe())

        except SyntaxError as e:
            st.error(f"代码语法错误: {str(e)}")
            st.info("""
常见语法错误：
1. 确保所有引号都是成对的
2. 检查是否缺少逗号或括号
3. 确保没有多余的空格或缩进
            """)
        except Exception as e:
            st.error("生成图表时出错:")
            st.code(traceback.format_exc())